import com.alibaba.druid.DbType;
import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLObject;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.expr.SQLQueryExpr;
import com.alibaba.druid.sql.ast.statement.*;

import java.util.Map;

/**
 * @author 支绍昆
 * 2024/3/20
 **/
public class Example {
    public static void main(String[] args) {
        String point = "a2d67fa7483449fb87cfe9c0d732bfa2";
//        single(point);
        batch(point);
    }

    private static void single(String point) {
        String sql = new Database().queryOne(point);
        System.out.println(sql);
        SQLStatement statement = SQLUtils.parseSingleStatement(sql, DbType.mysql);
        analysis(statement);
        System.out.println(statement.toString());
    }

    private static void batch(String point) {
        Database database = new Database();
        Map<String, String> sqlList = database.queryList(point);
        for (Map.Entry<String, String> entry : sqlList.entrySet()) {
            System.out.println("id : " + entry.getKey());
            SQLStatement statement = SQLUtils.parseSingleStatement(entry.getValue(), DbType.mysql);
            analysis(statement);
            System.out.println(statement.toString());
            System.out.println("====================================================================================================================================\n\n");
        }
    }

    private static void analysis(SQLObject sqlObject) {
        if (sqlObject instanceof SQLSelectStatement) {
            SQLSelectStatement selectStatement = (SQLSelectStatement) sqlObject;
            analysis(selectStatement.getSelect());
        } else if (sqlObject instanceof SQLSelect) {
            SQLSelect sqlSelect = (SQLSelect) sqlObject;
            analysis(sqlSelect.getQuery());
        } else if (sqlObject instanceof SQLSelectQueryBlock) {
            SQLSelectQueryBlock query = (SQLSelectQueryBlock) sqlObject;
            analysis(query.getFrom());
            analysis(query.getWhere());
        } else if (sqlObject instanceof SQLExprTableSource) {
            System.out.println("expr -> " + sqlObject);
            //获取上级select块
            SQLSelectQueryBlock parentSelect = findParentSelect(sqlObject);
            if (null == parentSelect) {
                return;
            }
            SQLExprTableSource expr = (SQLExprTableSource) sqlObject;
            //添加查询条件
            parentSelect.addWhere(new SQLIdentifierExpr((expr.getAlias() == null ? expr.toString() : expr.getAlias()) + ".depart_code" + " IN ($USER_DEPART$)"));
        } else if (sqlObject instanceof SQLJoinTableSource) {
            SQLJoinTableSource joinTableSource = (SQLJoinTableSource) sqlObject;
            analysis(joinTableSource.getLeft());
            analysis(joinTableSource.getRight());
        } else if (sqlObject instanceof SQLSubqueryTableSource) {
            SQLSubqueryTableSource subQueryTableSource = (SQLSubqueryTableSource) sqlObject;
            analysis(subQueryTableSource.getSelect());
        } else if (sqlObject instanceof SQLUnionQueryTableSource) {
            SQLUnionQueryTableSource queryTableSource = (SQLUnionQueryTableSource) sqlObject;
            analysis(queryTableSource.getUnion());
        } else if (sqlObject instanceof SQLUnionQuery) {
            SQLUnionQuery unionQuery = (SQLUnionQuery) sqlObject;
            unionQuery.getRelations().forEach(Example::analysis);
        } else if (sqlObject instanceof SQLBinaryOpExpr) {
            SQLBinaryOpExpr opExpr = (SQLBinaryOpExpr) sqlObject;
            analysis(opExpr.getLeft());
            analysis(opExpr.getRight());
        } else if (sqlObject instanceof SQLQueryExpr) {
            SQLQueryExpr queryExpr = (SQLQueryExpr) sqlObject;
            analysis(queryExpr.subQuery);
        }
    }

    private static SQLSelectQueryBlock findParentSelect(SQLObject sqlObject) {
        if (sqlObject instanceof SQLSelectQueryBlock) {
            return (SQLSelectQueryBlock) sqlObject;
        } else if (sqlObject instanceof SQLExprTableSource) {
            SQLExprTableSource expr = (SQLExprTableSource) sqlObject;
            return findParentSelect(expr.getParent());
        } else if (sqlObject instanceof SQLJoinTableSource) {
            SQLJoinTableSource join = (SQLJoinTableSource) sqlObject;
            return findParentSelect(join.getParent());
        }
        return null;
    }

}
